#include <unordered_map>
#include <unordered_set>
#include <bitset>
#include <gflags/gflags.h>
#include <glog/logging.h>
#include <fstream>
#include "core/grape.h"

using namespace std;
DEFINE_double(win,1,"label weight of incomingedages");
DEFINE_double(wout,1,"label weight of outgoingedges");

template<class VD, class ED, class MessageT, class ResultT, class TemporalT>
class LPWorker: public Worker<VD, ED, MessageT, ResultT, TemporalT> {
	void pEval(Fragment<VD, ED> &fragment, MessageBuffer<MessageT> &messageBuffer, ResultT &partialResult, TemporalT &temporal) {
		//read in weights of in and out edges from query
		double win = FLAGS_win;
		double wout = FLAGS_wout;
		
		unordered_map<VertexID, int> vertexlabel;
		for (auto vid : *fragment.getInnerVertices()){
			vertexlabel[vid] = vid;
		}

		//label update and creat messages
		for (auto v : vertexlabel){
			//label propagate for those vertices not at border
			if (!fragment.isBorderVertexWithOutgoingEdge(v.first) && !fragment.isBorderVertexWithIncomingEdge(v.first)){
				unordered_map<int, double> label_weight;
				//process outgoingedges
				for (auto e : fragment.getOutgoingEdges(v.first)){
					VertexID u = e.getDst();
					ED len = e.getAttr();
					if (label_weight.find(vertexlabel[u]) == label_weight.end()){
						label_weight[vertexlabel[u]] = wout * len;
					}else{
						label_weight[vertexlabel[u]] += wout *len;
					}
				}
				//process incomingedges
				for (auto e: fragment.getIncomingEdges(v.first)){
					VertexID u = e.getSrc();
					ED len = e.getAttr();
					if (label_weight.find(vertexlabel[u]) == label_weight.end()){
						label_weight[vertexlabel[u]] = win * len;
					}else{
						label_weight[vertexlabel[u]] += win * len;
					}
				}
	
				//decide new label for v
				double maxweight = 0;
				int newlabel = v.second;
				for (auto label : label_weight){
					if(maxweight < label.second){
						maxweight = label.second;
						newlabel = label.first;
					}
				}
				vertexlabel[v.first] = newlabel;
			}else if (fragment.isBorderVertexWithOutgoingEdge(v.first)){
				//create messages for those vertices with outgoingedges
				bitset<NUM_FRAGMENTS> dest = fragment.getMsgThroughOutgoingDest(v.first);
				for (int fid = 0; fid < dest.size(); fid ++){
					if (dest.test(fid)){
						messageBuffer.add_message(fid, PairMessage<int>(v.first,v.second));
					}
				}
			}else{
				//create messages for those vertices with incomingedges
				bitset<NUM_FRAGMENTS> dest = fragment.getMsgThroughIncomingDest(v.first);
				for (int fid = 0; fid < dest.size(); fid ++){
					if (dest.test(fid)){
						messageBuffer.add_message(fid, PairMessage<int>(v.first,v.second));
					}
				}
			}
		}
		partialResult = vertexlabel;
	}

	void incEval(Fragment<VD, ED> &fragment, MessageBuffer<MessageT> &messageBuffer, ResultT &partialResult, TemporalT &temporal) {
		double win  = FLAGS_win;
		double wout = FLAGS_wout;
		auto vertexlabel = partialResult;
		//recive messages
		unordered_map<VertexID, int> updated;
		for (auto item : messageBuffer.get_messages()){
			VertexID u = item.getID();
			int label = item.getValue();
			if (updated.find(u) == updated.end()){
				updated[u] = label;
			}
		}

		//update labels for vertices at border
		for (auto v : vertexlabel){
			unordered_map<int, double> label_weight;
			if (fragment.isBorderVertexWithOutgoingEdge(v.first) || fragment.isBorderVertexWithIncomingEdge(v.first)){
				for (auto e : fragment.getOutgoingEdges(v.first)){
					VertexID u = e.getDst();
					ED len  = e.getAttr();
					if (fragment.isInnerVertex(u)){
						if (label_weight.find(vertexlabel[u]) == label_weight.end()){
							label_weight[vertexlabel[u]] = wout * len;
						}else{
							label_weight[vertexlabel[u]] += wout *len;
						}
					}else{
						if (label_weight.find(updated[u]) == label_weight.end()){
							label_weight[updated[u]] = wout * len;
						}else{
							label_weight[updated[u]] += wout * len;
						}
					}
				}
				for (auto e : fragment.getIncomingEdges(v.first)){
					VertexID u = e.getSrc();
					ED len = e.getAttr();
					if (fragment.isInnerVertex(u)){
						if (label_weight.find(vertexlabel[u]) == label_weight.end()){
							label_weight[vertexlabel[u]] = win * len;
						}else{
							label_weight[vertexlabel[u]] += win * len;
						}
					}else{
						if (label_weight.find(updated[u]) == label_weight.end()){
							label_weight[updated[u]] = win * len;
						}else{
							label_weight[updated[u]] += win * len;
						}
					}
				}
			}

			//decide new label
			double maxweight = 0;
			int newlabel = v.second;
			for (auto label : label_weight){
				if (maxweight < label.second){
					maxweight = label.second;
					newlabel = label.first;
				}
			}
			vertexlabel[v.first] = newlabel;
		}
		//check to stop
		int finalvertices = 0;
		for (auto v : vertexlabel){
			unordered_map<int, double> label_weight;
			for (auto e : fragment.getOutgoingEdges(v.first)){
				VertexID u = e.getDst();
				ED len = e.getAttr();
				if(fragment.isInnerVertex(u)){
					if(label_weight.find(vertexlabel[u]) == label_weight.end()){
						label_weight[vertexlabel[u]] = wout * len;
					}else{
						label_weight[vertexlabel[u]] += wout * len;
					}
				}else{
					if(label_weight.find(updated[u]) == label_weight.end()){
						label_weight[updated[u]] = wout * len;
					}else{
						label_weight[updated[u]] += wout * len;
					}
				}
			}
			for (auto e : fragment.getIncomingEdges(v.first)){
				VertexID u = e.getSrc();
				ED len = e.getAttr();
				if(fragment.isInnerVertex(u)){
					if(label_weight.find(vertexlabel[u]) == label_weight.end()){
						label_weight[vertexlabel[u]] = win * len;
					}else{
						label_weight[vertexlabel[u]] += win * len;
					}
				}else{
					if(label_weight.find(updated[u]) == label_weight.end()){
						label_weight[updated[u]] = win * len;
					}else{
						label_weight[updated[u]] += win * len;
					}
				}
			}

			//whether fianl label is the one have maxweight
			double maxweight = 0;
			int usedlabel = v.second;
			for (auto label : label_weight){
				if (maxweight < label.second){
					maxweight = label.second;
					usedlabel = label.first;
				}
			}
			if (v.second == usedlabel){
				finalvertices += 1;
				continue;
			}else{
				break;
			}
		}
		
		//if not stop, update labels for vertices not at border and create messages
		if (finalvertices != vertexlabel.size()){
			for (auto v : vertexlabel){
				//label propagate for those vertices not at border
				if (!fragment.isBorderVertexWithOutgoingEdge(v.first) && !fragment.isBorderVertexWithIncomingEdge(v.first)){
					unordered_map<int, double> label_weight;
					//process outgoingedges
					for (auto e : fragment.getOutgoingEdges(v.first)){
						VertexID u = e.getDst();
						ED len = e.getAttr();
						if (label_weight.find(vertexlabel[u]) == label_weight.end()){
							label_weight[vertexlabel[u]] = wout * len;
						}else{
							label_weight[vertexlabel[u]] += wout *len;
						}
					}
					//process incomingedges
					for (auto e: fragment.getIncomingEdges(v.first)){
						VertexID u = e.getSrc();
						ED len = e.getAttr();
						if (label_weight.find(vertexlabel[u]) == label_weight.end()){
							label_weight[vertexlabel[u]] = win * len;
						}else{
							label_weight[vertexlabel[u]] += win * len;
						}
					}

					//decide new label for v
					double maxweight = 0;
					int newlabel = v.second;
					for (auto label : label_weight){
						if(maxweight < label.second){
							maxweight = label.second;
							newlabel = label.first;
						}
					}
					vertexlabel[v.first] = newlabel;
				}else if (fragment.isBorderVertexWithOutgoingEdge(v.first)){
					//create messages for those vertices with outgoingedges
					bitset<NUM_FRAGMENTS> dest = fragment.getMsgThroughOutgoingDest(v.first);
					for (int fid = 0; fid < dest.size(); fid ++){
						if (dest.test(fid)){
							messageBuffer.add_message(fid, PairMessage<int>(v.first,v.second));
						}
					}
				}else{
					//create messages for those vertices with incomingedges
					bitset<NUM_FRAGMENTS> dest = fragment.getMsgThroughIncomingDest(v.first);
					for (int fid = 0; fid < dest.size(); fid ++){
						if (dest.test(fid)){
							messageBuffer.add_message(fid, PairMessage<int>(v.first,v.second));
						}
					}
				}
			}
		}
		partialResult = vertexlabel;
	}

	void assemble(vector<ResultT> &partialResults) {
		map<VertexID, int> finalResult;
		unordered_map<VertexID, int> labels;
		for (auto result : partialResults) {
			for (auto item : result) {
				VertexID u = item.first;
				int label = item.second;
				finalResult[u] = label;
				if(labels.find(label) == labels.end()){
					labels[label] = label;
				}
			}
		}

		if (this->result_file != "") {
			std::ofstream fout;
			fout.open(this->result_file);
			fout << "the number of commities is " << labels.size() << endl;
			for (auto item : finalResult) {
				fout << item.first << '\t' << item.second << endl;
			}
			fout.close();
		}
	}	
};

int main(int argc, char *argv[]) {

	FLAGS_stderrthreshold = 0;

	google::SetUsageMessage("Usage: mpiexec.hydra [mpi_opts] ./lp [grape_opts]");
	google::ParseCommandLineFlags(&argc, &argv, true);
	init_workers();

	WorkerParams param;

	param.algorithm = "lp";
	param.enable_assemble = true;
	param.graph_load_strategy = BOTH_OUT_IN;
	param.use_pattern = false;

	google::ShutDownCommandLineFlags();

	LPWorker<int, int, PairMessage<int>, unordered_map<VertexID, int>, VertexID> lpWorker;

	google::InitGoogleLogging("lp");
	lpWorker.run(param);
	google::ShutdownGoogleLogging();

	worker_finalize();
	return 0;
}
